/* nbd-server.c */

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <errno.h>

#include <time.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>

#include "nbd-server.h"
#include "ip.h"
#include "tcp.h"
#include "dump.h"
#include "pack.h"
#include "nbd-dev.h"

#define CONN_TIMEOUT  300

#define NBD_MAGIC_INIT       0x0000420281861253llu
#define NBD_MAGIC_REQUEST    0x25609513u
#define NBD_MAGIC_REPLY      0x67446698u
#define NBD_MAGIC_EXT_REPLY  0xb8844f80u

#define REQ_STATUS_INACTIVE  0   /* no request being processed (must be 0) */
#define REQ_STATUS_HEADER    1   /* sending reply header */
#define REQ_STATUS_DATA      2   /* sending reply data */

typedef struct NBD_REQ {
  unsigned long long handle;
  unsigned long long off;
  unsigned int len;
  NBD_DEV *dev;
} NBD_REQ;


/**
 * Read a NBD request, return -1 on error.
 */
static int read_nbd_request(int sock, NBD_REQ *req)
{
  unsigned char buf[28];
  int len;

  dump_msg("-> reading request\n");

  while (42) {
    len = tcp_recv(sock, buf, sizeof(buf));
    if (len != 0 || (len == 0 && errno != EINTR))
      break;
  }
  if (len < 0) {
    if (errno == ECONNRESET)
      dump_msg("** TCP connection reset by peer\n");
    else
      dump_msg("** TCP read error (errno=%d)\n", errno);
    return -1;
  }
  if (len == 0) {
    dump_msg("** Connection closed while reading request\n");
    return -1;
  }
  if (len != 28) {
    dump_msg("** Unexpected request length: %d (expecting 28)\n", len);
    return -1;
  }

  if (unpack_u32_be(buf) != NBD_MAGIC_REQUEST) {
    dump_msg("** Bad request MAGIC (%08x)\n", unpack_u32_be(buf));
    return -1;
  }

  if (unpack_u32_be(buf + 4) != 0) {
    dump_msg("** Bad request TYPE: %d\n", unpack_u32_be(buf + 4));
    return -1;
  }

  req->handle = unpack_u64_be(buf + 8);
  req->off = unpack_u64_be(buf + 16);
  req->len = unpack_u32_be(buf + 24);

  dump_msg("  -> handle=%016llx\n", req->handle);
  dump_msg("  -> off=%016llx\n", req->off);
  dump_msg("  -> len=%08x\n", req->len);

  return 0;
}

/**
 * Send NBD reply, return -1 on error.
 */
static int write_nbd_response(int sock, NBD_REQ *req)
{
  int pos, len;
  unsigned char header[24];

  // -- build header
  // this is the standard NBD reply header:
  pack_u32_be(header, NBD_MAGIC_REPLY);
  pack_u32_be(header + 4, 0);
  pack_u64_be(header + 8, req->handle);

  // these are extensions:
  pack_u32_be(header + 16, NBD_MAGIC_EXT_REPLY);
  pack_u32_be(header + 20, 0);

  // -- write header
  dump_msg("-> write header\n");
  len = tcp_send(sock, header, 24);
  dump_msg("  -> length=%d\n", len);
  if (len <= 0) {
    dump_msg("** Error sending reply (errno=%d)\n", errno);
    return -1;
  }

  // -- write data
  pos = 0;
  while (pos < req->len) {
    //char buf[1460];
    char buf[65536];
    int dev_len;

    len = req->len - pos;
    len = (len > sizeof(buf)) ? sizeof(buf) : len;

    // read device data
    dev_len = nbd_dev_read_data(req->dev, buf, len, req->off + pos);
    if (dev_len <= 0) {
      dump_msg("** Error reading device\n");
      return -1;
    }
    if (dev_len != len) {
      dump_msg("** Unexpected data length returned by device\n");
      return -1;
    }

    // write response
    dump_msg("-> writing data (off=%llu,pos=%u/%u,len=%u)\n", req->off, pos, req->len, len);
    len = tcp_send(sock, buf, len);
    dump_msg("  -> length=%d\n", len);
    if (len <= 0) {
      dump_msg("** Error sending reply (errno=%d)\n", errno);
      return -1;
    }
    pos += len;
  }

  //usleep(10000);
  return 0;
}

/**
 * Send initial response to NBD client
 */
static int write_nbd_init(int sock, NBD_DEV *dev)
{
  unsigned char header[152];
  long long dev_size;
  int len;

  dev_size = nbd_dev_get_size(dev);
  if (dev_size < 0) {
    dump_msg("** Can't get device size\n");
    return -1;
  }

  memcpy(header, "NBDMAGIC", 8);
  pack_u64_be(header + 8, NBD_MAGIC_INIT);
  pack_u64_be(header + 16, dev_size);
  pack_u32_be(header + 24, 3);   // flags
  memset(header + 28, 0, 124);

  dump_data(0, header, 24);

  dump_msg("-> sending init packet\n");
  len = tcp_send(sock, header, 152);
  if (len <= 0) {
    dump_msg("ERROR can't sending init header\n");
    return -1;
  }
  return 0;
}

/**
 * Handle a NBD connection.
 */
void nbd_server_handle_connection(int sock, const char *image_file)
{
  NBD_REQ req;
  NBD_DEV dev;

  // open disk image
  if (nbd_dev_open_image(&dev, image_file) < 0) {
    dump_msg("ERROR: can't open image file '%s'\n", image_file);
    close(sock);
    return;
  }

  // init request
  memset(&req, 0, sizeof(NBD_REQ));
  req.dev = &dev;

  // send INIT NBD header
  if (write_nbd_init(sock, req.dev) != 0)
    return;

  // adjust socket options
  //tcp_disable_delay(sock);

  // server loop
  while (42) {
    dump_msg("-> loop\n");

    // read NBD request
    if (read_nbd_request(sock, &req) < 0) {
      dump_msg("-> closing (read)\n");
      tcp_close(sock);
      return;
    }
    
    // write NBD response
    if (write_nbd_response(sock, &req) < 0) {
      dump_msg("-> closing (write)\n");
      tcp_close(sock);
      return;
    }
  }
}

/**
 * Open a NBD server.
 */
int nbd_server_open(NBD_SERVER *srv)
{
  srv->sock = tcp_open_server(0, 31013, 1);
  if (srv->sock < 0) {
    dump_msg("Error: can't create NBD server on port 31013.\n");
    return -1;
  }
  return 0;
}

/**
 * Close a NBD server.
 */
void nbd_server_close(NBD_SERVER *srv)
{
  tcp_close(srv->sock);
}

/**
 * Run a NBD server.
 */
void nbd_server_run(NBD_SERVER *srv, char *image_file)
{
  net_u32 ip;
  int port;

  while (42) {
    pid_t pid;
    int fd = tcp_accept(srv->sock, &ip, &port);
    dump_msg("-> got connection from %d.%d.%d.%d:%d\n",
	     (ip>>24)&0xff, (ip>>16)&0xff, (ip>>8)&0xff, ip&0xff, port);
    pid = fork();
    if (pid == 0) {
      nbd_server_close(srv);
      nbd_server_handle_connection(fd, image_file);
      _exit(0);
    }
    if (pid < 0)
      dump_msg("** ERROR: can't fork()\n");
  }
}
